package websocket

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

// WebSocketServer WebSocket服务器
type WebSocketServer struct {
	upgrader    websocket.Upgrader
	clients     map[*websocket.Conn]*Client
	clientsMux  sync.RWMutex
	broadcast   chan []byte
	register    chan *Client
	unregister  chan *Client
	ctx         context.Context
	cancel      context.CancelFunc
}

// Client 客户端连接
type Client struct {
	conn     *websocket.Conn
	send     chan []byte
	server   *WebSocketServer
	id       string
	lastSeen time.Time
}

// Message WebSocket消息结构
type Message struct {
	Type      string      `json:"type"`
	Content   string      `json:"content,omitempty"`
	Data      interface{} `json:"data,omitempty"`
	Timestamp time.Time   `json:"timestamp"`
	ClientID  string      `json:"client_id,omitempty"`
}

// Response WebSocket响应结构
type Response struct {
	Type      string      `json:"type"`
	Success   bool        `json:"success"`
	Message   string      `json:"message,omitempty"`
	Data      interface{} `json:"data,omitempty"`
	Timestamp time.Time   `json:"timestamp"`
}

// NewWebSocketServer 创建新的WebSocket服务器
func NewWebSocketServer() *WebSocketServer {
	ctx, cancel := context.WithCancel(context.Background())
	
	return &WebSocketServer{
		upgrader: websocket.Upgrader{
			CheckOrigin: func(r *http.Request) bool {
				// 允许所有来源，生产环境中应该限制
				return true
			},
			ReadBufferSize:  1024,
			WriteBufferSize: 1024,
		},
		clients:    make(map[*websocket.Conn]*Client),
		broadcast:  make(chan []byte),
		register:   make(chan *Client),
		unregister: make(chan *Client),
		ctx:        ctx,
		cancel:     cancel,
	}
}

// Start 启动WebSocket服务器
func (s *WebSocketServer) Start(port string) error {
	// 启动消息处理协程
	go s.handleMessages()
	
	// 设置WebSocket路由
	http.HandleFunc("/ws", s.handleWebSocket)
	
	// 设置静态文件服务（可选）
	http.HandleFunc("/", s.handleIndex)
	
	log.Printf("WebSocket服务器启动在端口 %s", port)
	
	// 启动HTTP服务器
	return http.ListenAndServe(":"+port, nil)
}

// Stop 停止WebSocket服务器
func (s *WebSocketServer) Stop() {
	s.cancel()
	
	// 关闭所有客户端连接
	s.clientsMux.Lock()
	for conn, client := range s.clients {
		close(client.send)
		conn.Close()
		delete(s.clients, conn)
	}
	s.clientsMux.Unlock()
	
	log.Println("WebSocket服务器已停止")
}

// handleWebSocket 处理WebSocket连接
func (s *WebSocketServer) handleWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := s.upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket升级失败: %v", err)
		return
	}
	
	// 创建客户端
	client := &Client{
		conn:     conn,
		send:     make(chan []byte, 256),
		server:   s,
		id:       generateClientID(),
		lastSeen: time.Now(),
	}
	
	// 注册客户端
	s.register <- client
	
	// 启动客户端处理协程
	go client.writePump()
	go client.readPump()
	
	log.Printf("新客户端连接: %s", client.id)
}

// handleIndex 处理首页请求
func (s *WebSocketServer) handleIndex(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" {
		http.Error(w, "Not found", http.StatusNotFound)
		return
	}
	
	// 返回简单的WebSocket测试页面
	html := `
<!DOCTYPE html>
<html>
<head>
    <title>Loomi WebSocket Test</title>
    <meta charset="utf-8">
</head>
<body>
    <h1>Loomi WebSocket 测试页面</h1>
    <div id="messages"></div>
    <input type="text" id="messageInput" placeholder="输入消息...">
    <button onclick="sendMessage()">发送</button>
    
    <script>
        const ws = new WebSocket('ws://localhost:8080/ws');
        const messages = document.getElementById('messages');
        
        ws.onopen = function() {
            addMessage('连接已建立');
        };
        
        ws.onmessage = function(event) {
            const data = JSON.parse(event.data);
            addMessage('收到: ' + JSON.stringify(data, null, 2));
        };
        
        ws.onclose = function() {
            addMessage('连接已关闭');
        };
        
        function sendMessage() {
            const input = document.getElementById('messageInput');
            const message = {
                type: 'chat',
                content: input.value,
                timestamp: new Date().toISOString()
            };
            ws.send(JSON.stringify(message));
            input.value = '';
        }
        
        function addMessage(msg) {
            const div = document.createElement('div');
            div.innerHTML = '<pre>' + msg + '</pre>';
            messages.appendChild(div);
            messages.scrollTop = messages.scrollHeight;
        }
        
        document.getElementById('messageInput').addEventListener('keypress', function(e) {
            if (e.key === 'Enter') {
                sendMessage();
            }
        });
    </script>
</body>
</html>`
	
	w.Header().Set("Content-Type", "text/html")
	w.Write([]byte(html))
}

// handleMessages 处理消息分发
func (s *WebSocketServer) handleMessages() {
	for {
		select {
		case client := <-s.register:
			s.clientsMux.Lock()
			s.clients[client.conn] = client
			s.clientsMux.Unlock()
			
			// 发送欢迎消息
			welcome := Response{
				Type:      "welcome",
				Success:   true,
				Message:   "连接成功",
				Data:      map[string]string{"client_id": client.id},
				Timestamp: time.Now(),
			}
			client.sendResponse(welcome)
			
		case client := <-s.unregister:
			s.clientsMux.Lock()
			if _, ok := s.clients[client.conn]; ok {
				delete(s.clients, client.conn)
				close(client.send)
			}
			s.clientsMux.Unlock()
			log.Printf("客户端断开连接: %s", client.id)
			
		case message := <-s.broadcast:
			s.clientsMux.RLock()
			for conn, client := range s.clients {
				select {
				case client.send <- message:
				default:
					close(client.send)
					delete(s.clients, conn)
				}
			}
			s.clientsMux.RUnlock()
			
		case <-s.ctx.Done():
			return
		}
	}
}

// Broadcast 广播消息给所有客户端
func (s *WebSocketServer) Broadcast(message interface{}) {
	data, err := json.Marshal(message)
	if err != nil {
		log.Printf("消息序列化失败: %v", err)
		return
	}
	
	select {
	case s.broadcast <- data:
	default:
		log.Println("广播通道已满，消息丢弃")
	}
}

// GetConnectedClients 获取连接的客户端数量
func (s *WebSocketServer) GetConnectedClients() int {
	s.clientsMux.RLock()
	defer s.clientsMux.RUnlock()
	return len(s.clients)
}

// generateClientID 生成客户端ID
func generateClientID() string {
	return fmt.Sprintf("client_%d", time.Now().UnixNano())
}